package com.mrli.websocket;

import com.mrli.openai.OpenAiService;
import com.mrli.openai.completion.chat.ChatCompletionChoice;
import com.mrli.openai.completion.chat.ChatCompletionRequest;
import com.mrli.openai.completion.chat.ChatCompletionResult;
import com.mrli.openai.completion.chat.ChatMessage;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.io.IOException;
import java.math.BigInteger;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * description
 *
 * @author LiDaShuai
 * @version 1.0
 * @project mrli-mall-auto-issu
 * @date: 2023/2/8 16:46
 * @copyright 2009–2022xxxxx all rights reserved.
 */
@Component
@Slf4j
public class MessageHandler extends TextWebSocketHandler {

    private static final Map<BigInteger, WebSocketSession> SESSIONS = new ConcurrentHashMap<>();
    private final OpenAiService service;


    public MessageHandler(OpenAiService service) {
        this.service = service;
    }
//
//    @Autowired
//    private MessageDAO messageDAO;

    /**
     * 客户端与服务端建立websocket连接后执行
     *
     * @param session
     * @throws Exception
     */
    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        Object uid = session.getAttributes().get("userId");
        if (uid == null) {
            session.sendMessage(new TextMessage("建立链接失败"));
            session.close();
            return;
        }
        BigInteger userId = (BigInteger) uid;
        log.info("WebSocket建立连接中,连接用户ID：{}", userId);
        WebSocketSession webSocketSession = SESSIONS.get(userId);
        /**
         * webSocketSession不为空,说明已经有人登陆账号,应该删除登陆的WebSocket对象
         */
        if (webSocketSession != null) {
            webSocketSession.close();
        }
        SESSIONS.put(userId, session);
        log.info("建立连接完成,当前在线人数为：{}", SESSIONS.size());
    }

    /**
     * 服务端接收到消息后执行
     *
     * @param session
     * @param textMessage 消息体
     */
    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage textMessage) {
        Object userId1 = session.getAttributes().get("userId");
        if (userId1 == null) {
            log.error("获取用户信息失败");
            return;
        }
        BigInteger userId = (BigInteger) userId1;
        String payload = textMessage.getPayload();
        if ("1".equals(payload)) {
            return;
        }
        WebSocketSession toSession = SESSIONS.get(userId);
        if (toSession != null && toSession.isOpen()) {
            List<ChatMessage> messages = new ArrayList<>();
            ChatMessage e = new ChatMessage();
            e.setRole("user");
            e.setContent(textMessage.getPayload());
            messages.add(e);
            ChatCompletionRequest completionRequest = ChatCompletionRequest.builder()
                    .model("gpt-3.5-turbo")
                    .messages(messages)
                    .user("李大帅3")
                    .n(1)
                    .temperature(0d)
                    .topP(1D)
                    .build();
            try {
                ChatCompletionResult completion = service.createChatCompletion(completionRequest);
                List<ChatCompletionChoice> choices = completion.getChoices();
                choices.forEach(s -> {
                    try {
                        toSession.sendMessage(new TextMessage(s.getMessage().getContent()));
                    } catch (IOException ioException) {
                        ioException.printStackTrace();
                        throw new RuntimeException();
                    }
                });
            } catch (Exception e1) {
                e1.printStackTrace();
                try {
                    toSession.sendMessage(new TextMessage("系统开小差了"));
                } catch (IOException ioException) {
                    ioException.printStackTrace();
                }
            }


        }
        log.info("收到客户端发来的消息：{}，用户id：{}", textMessage.getPayload(), userId);
    }

    /**
     * 客户端与服务端关闭websocket连接后执行
     *
     * @param session
     * @param status
     */
    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) {
        Object userId1 = session.getAttributes().get("userId");
        if (userId1 == null) {
            return;
        }
        BigInteger userId = new BigInteger(userId1.toString());
        WebSocketSession webSocketSession = SESSIONS.get(userId);
        /**
         * webSocketSession,说明已经有人登陆账号,应该删除登陆的WebSocket对象
         */
        if (webSocketSession != null) {
            try {
                webSocketSession.close();
            } catch (IOException e) {
                e.printStackTrace();
                log.error("用户:{},关闭session失败", userId);
            }
            SESSIONS.remove(userId);
            log.info("连接断开,当前在线人数为：{}", SESSIONS.size());
        }
    }

    /**
     * 推送消息到指定用户
     *
     * @param userId  用户ID
     * @param message 发送的消息
     */
    public static void sendMessageByUser(BigInteger userId, String message) {
        log.info("用户ID：" + userId + ",推送内容：" + message);
        WebSocketSession session = SESSIONS.get(userId);
        try {
            if (session != null) {
                session.sendMessage(new TextMessage(message));
            }
        } catch (IOException e) {
            log.error("推送消息到指定用户发生错误：" + e.getMessage(), e);
        }
    }

    /**
     * 群发消息
     *
     * @param message 发送的消息
     */
    public static void sendAllMessage(String message) {
        log.info("发送消息：{}", message);
        SESSIONS.values().forEach(webSocketSession -> {
            try {
                webSocketSession.sendMessage(new TextMessage(message));
            } catch (IOException e) {
                log.error("群发消息发生错误：" + e.getMessage(), e);
            }
        });
    }
}